import sys
import numpy as np

# 针对float32类型的误差容忍度
relative_tol = 1e-5
absolute_tol = 1e-6
error_tol = 1e-3  # 允许的误差比例上限

def verify_single_result(output_path, golden_path):
    """验证单个输出文件与参考文件的一致性"""
    # 读取数据（float32类型，与main.cpp中的float对应）
    output = np.fromfile(output_path, dtype=np.float32).reshape(-1)
    golden = np.fromfile(golden_path, dtype=np.float32).reshape(-1)

    # 检查形状是否一致
    if output.shape != golden.shape:
        raise ValueError(f"形状不匹配: 输出{output.shape} vs 参考{golden.shape}")

    # 计算误差
    match = np.isclose(output, golden, rtol=relative_tol, atol=absolute_tol, equal_nan=True)
    mismatch_indices = np.where(~match)[0]

    # 打印前100个不匹配项
    for i, idx in enumerate(mismatch_indices[:10]):
        golden_val = golden[idx]
        output_val = output[idx]
        rdiff = abs(output_val - golden_val) / (abs(golden_val) + 1e-10)  # 避免除零
        print(f"数据索引: {idx:06d}, 预期: {golden_val:.6f}, 实际: {output_val:.6f}, 相对误差: {rdiff:.6f}")

    # 计算误差比例
    error_ratio = len(mismatch_indices) / len(golden)
    print(f"误差比例: {error_ratio:.4f}, 容忍上限: {error_tol:.4f}")
    return error_ratio <= error_tol

if __name__ == '__main__':
    try:
        # 验证两个输出：updated_mean和updated_var
        if len(sys.argv) != 5:
            raise ValueError("参数错误: 请输入 [updated_mean输出] [updated_mean参考] [updated_var输出] [updated_var参考]")
        
        mean_ok = verify_single_result(sys.argv[1], sys.argv[2])
        var_ok = verify_single_result(sys.argv[3], sys.argv[4])

        if mean_ok and var_ok:
            print("测试通过")
        else:
            raise ValueError("[ERROR] 结果验证失败")
    except Exception as e:
        print(e)
        sys.exit(1)